
# 包的帮助与介绍
import importlib
import inspect
import sys
from typing import List, Dict, Any, Optional, Union


def query_package(package_name: str, submodule_path: str = None, max_depth: int = 2, current_depth: int = 0) -> Dict[str, Any]:
    """
    递归查询包及其子模块的信息
    
    Args:
        package_name: 包名称 (如 'sklearn')
        submodule_path: 子模块路径 (如 'linear_model')
        max_depth: 最大递归深度
        current_depth: 当前递归深度
    
    Returns:
        包含模块信息的字典
    """
    result = {}
    
    # 构建完整的模块路径
    full_path = package_name
    if submodule_path:
        full_path = f"{package_name}.{submodule_path}"
    
    # 尝试导入模块
    try:
        module = importlib.import_module(full_path)
        
        # 获取基本信息
        result["name"] = full_path
        result["dir"] = dir(module)
        
        # 如果达到最大深度，则不再递归
        if current_depth >= max_depth:
            return result
        
        # 获取子模块和类
        submodules = {}
        classes = {}
        
        for item_name in result["dir"]:
            # 跳过私有属性
            if item_name.startswith("_"):
                continue
            
            try:
                item = getattr(module, item_name)
                
                # 检查是否为模块
                if inspect.ismodule(item):
                    # 检查该模块是否属于当前包
                    if hasattr(item, "__name__") and item.__name__.startswith(package_name):
                        submodule_info = query_package(
                            package_name,
                            f"{submodule_path}.{item_name}" if submodule_path else item_name,
                            max_depth,
                            current_depth + 1
                        )
                        submodules[item_name] = submodule_info
                
                # 检查是否为类
                elif inspect.isclass(item):
                    classes[item_name] = {
                        "name": item_name,
                        "dir": dir(item)
                    }
            except (ImportError, AttributeError):
                continue
        
        if submodules:
            result["submodules"] = submodules
        if classes:
            result["classes"] = classes
            
    except ImportError as e:
        result["error"] = f"无法导入模块: {e}"
    
    return result


def search_package(package_name: str, query: str, submodule_path: str = None, max_depth: int = 2) -> List[str]:
    """
    在包中搜索匹配查询的项
    
    Args:
        package_name: 包名称 (如 'sklearn')
        query: 搜索关键词 (如 'regression')
        submodule_path: 子模块路径 (可选)
        max_depth: 最大递归深度
    
    Returns:
        匹配项的列表
    """
    results = []
    package_info = query_package(package_name, submodule_path, max_depth)
    
    def search_in_dict(info_dict: Dict[str, Any], path: str = ""):
        # 搜索当前模块的 dir 列表
        if "dir" in info_dict:
            for item in info_dict["dir"]:
                if query.lower() in item.lower():
                    item_path = f"{path}.{item}" if path else item
                    results.append(f"{item_path} (in {info_dict.get('name', path)})")
        
        # 搜索子模块
        if "submodules" in info_dict:
            for submodule_name, submodule_info in info_dict["submodules"].items():
                submodule_path = f"{path}.{submodule_name}" if path else submodule_name
                search_in_dict(submodule_info, submodule_path)
        
        # 搜索类
        if "classes" in info_dict:
            for class_name, class_info in info_dict["classes"].items():
                # 类名匹配
                if query.lower() in class_name.lower():
                    class_path = f"{path}.{class_name}" if path else class_name
                    results.append(f"{class_path} (class in {info_dict.get('name', path)})")
                
                # 类的方法和属性匹配
                for item in class_info.get("dir", []):
                    if query.lower() in item.lower():
                        item_path = f"{path}.{class_name}.{item}" if path else f"{class_name}.{item}"
                        results.append(f"{item_path} (in class {class_name})")
    
    search_in_dict(package_info, submodule_path or package_name)
    return results


def print_module_structure(package_name: str, submodule_path: str = None, max_depth: int = 1):
    """
    打印模块结构
    
    Args:
        package_name: 包名称 (如 'sklearn')
        submodule_path: 子模块路径 (可选)
        max_depth: 最大递归深度(注：第一个子模块为深度1)
    """
    package_info = query_package(package_name, submodule_path, max_depth)
    
    def print_dict(info_dict: Dict[str, Any], indent: int = 0):
        if "name" in info_dict:
            print("  " * indent + f"📦 {info_dict['name']}")
        
        if "classes" in info_dict:
            for class_name in sorted(info_dict["classes"].keys()):
                print("  " * (indent + 1) + f"🔹 {class_name} (类)")
        
        if "submodules" in info_dict:
            for submodule_name, submodule_info in sorted(info_dict["submodules"].items()):
                print_dict(submodule_info, indent + 1)
    
    print_dict(package_info)


# 示例用法
if __name__ == "__main__":
    # 1. 查询 sklearn 包结构
    print("\n===== sklearn 包结构 =====")
    print_module_structure("sklearn", max_depth=1)
    
    # 2. 查询 sklearn.linear_model 子模块
    print("\n===== sklearn.linear_model 子模块结构 =====")
    print_module_structure("sklearn", "linear_model", max_depth=1)
    
    # 3. 搜索包含 "regression" 的项
    print("\n===== 搜索包含 'regression' 的项 =====")
    results = search_package("sklearn", "regression", max_depth=2)
    for result in results[:20]:  # 限制输出数量
        print(result)
    
    if len(results) > 20:
        print(f"... 还有 {len(results) - 20} 个结果")
    
    # 4. 交互式查询
    print("\n===== 交互式查询 =====")
    while True:
        print("\n选择操作:")
        print("1. 查看包结构")
        print("2. 搜索关键词")
        print("3. 退出")
        
        choice = input("请输入选项 (1/2/3): ").strip()
        
        if choice == "1":
            pkg = input("请输入包名 (例如 'sklearn'): ").strip()
            submodule = input("请输入子模块路径 (可选): ").strip() or None
            depth = input("请输入最大深度 (默认1): ").strip() or "1"
            try:
                print_module_structure(pkg, submodule, int(depth))
            except Exception as e:
                print(f"错误: {e}")
                
        elif choice == "2":
            pkg = input("请输入包名 (例如 'sklearn'): ").strip()
            query_str = input("请输入搜索关键词: ").strip()
            submodule = input("请输入子模块路径 (可选): ").strip() or None
            depth = input("请输入最大深度 (默认2): ").strip() or "2"
            try:
                results = search_package(pkg, query_str, submodule, int(depth))
                for result in results[:30]:
                    print(result)
                if len(results) > 30:
                    print(f"... 还有 {len(results) - 30} 个结果")
            except Exception as e:
                print(f"错误: {e}")
                
        elif choice == "3":
            break
        else:
            print("无效选项，请重新输入")
